#!/usr/bin/env python3

from glob import glob
import json
import os
import sys
wrf_script_root = os.path.dirname(os.path.realpath(__file__)) + '/..'
sys.path.append(f'{wrf_script_root}/utils')
sys.path.append(f'{wrf_script_root}/operators')
sys.path.append(f'{wrf_script_root}/osse/sim_obs')
import numpy as np
from pprint import pprint
from jinja2 import Template
import wrf_operators as wrf
from sim_obs import sim_synop, sim_sound, sim_profiler
from utils import cmd_args, cli, parse_config, run, submit_job

args = cmd_args()

config = parse_config(args.config_json)
config['custom']['littler'] = {
	'dir_pattern': args.work_root + '/sim_obs',
	'file_pattern': 'ob.{{time.format("YYYYMMDDHHmm")}}.txt'
}
args.littler_root = ''

# Simulate observation.
def sim_obs(time):
	time_str = time.format("YYYYMMDDHHmm")
	sim_obs_dir = args.work_root + '/sim_obs'
	sim_littler = f'{sim_obs_dir}/ob.{time_str}.txt'
	if not os.path.isfile(sim_littler):
		# SYNOP
		synop_littler = f'{sim_obs_dir}/synop.{time.format("YYYYMMDDHHmm")}.txt'
		if not os.path.isfile(synop_littler):
			exe = f'{wrf_script_root}/osse/sim_obs/sim_synop.py'
			submit_job(f'{exe} -o {synop_littler} --nature-run-root {args.nature_run_root} --start-time {time_str} --end-time {time_str} -j {args.config_json} -n 20', 1, config, args, cores=20, mpi=False)
		# SOUND
		sound_littler = f'{sim_obs_dir}/sound.{time.format("YYYYMMDDHHmm")}.txt'
		if not os.path.isfile(sound_littler):
			exe = f'{wrf_script_root}/osse/sim_obs/sim_sound.py'
			submit_job(f'{exe} -o {sound_littler} --nature-run-root {args.nature_run_root} --start-time {time_str} --end-time {time_str} -j {args.config_json} -n 20', 1, config, args, cores=20, mpi=False)
		# PROFILER
		profiler_littler = f'{sim_obs_dir}/profiler.{time.format("YYYYMMDDHHmm")}.txt'
		if not os.path.isfile(profiler_littler):
			exe = f'{wrf_script_root}/osse/sim_obs/sim_profiler.py'
			submit_job(f'{exe} -o {profiler_littler} --nature-run-root {args.nature_run_root} --start-time {time_str} --end-time {time_str} -j {args.config_json} -n 20', 1, config, args, cores=20, mpi=False)
		run(f'cat {synop_littler} {sound_littler} {profiler_littler} > {sim_littler}')

def run_wps(config):
	wrf.config_wps(args.work_root, args.wps_root, args.geog_root, config, args)
	wrf.run_wps_geogrid(args.work_root, args.wps_root, config, args)
	wrf.run_wps_ungrib_metgrid(args.work_root, args.wps_root, args.bkg_root, config, args)

def run_spinup(config):
	wrf.config_wrf(args.work_root, args.wrf_root, config, args, tag='spinup')
	wrf.run_real(args.work_root, args.wrf_root, config, args, tag='spinup')
	wrf.run_wrf(args.work_root, args.wrf_root, config, args, tag='spinup')

def run_osse(exp):
	wrf.config_wrf(args.work_root, args.wrf_root, config, args, tag=exp)
	wrf.run_real(args.work_root, args.wrf_root, config, args, tag=exp)

	if exp == 'sens':
		wrf.config_wrfda(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_obsproc(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_3dvar(args.work_root, args.wrfda_root, config, args, tag=exp)
		wrf.run_wrfda_update_bc(args.work_root, args.wrfda_root, config, args, tag=exp)

	wrf.run_wrf(args.work_root, args.wrf_root, config, args, tag=exp)
	wrf.run_upp(args.work_root, args.upp_root, config, args, tag=exp)
	wrf.run_met(args.work_root, args.met_root, config, args, tag=exp)

sim_obs(config['custom']['start_time'])
sim_obs(config['custom']['end_time'])

spinup_hours = get_value(config['custom'], 'spinup_hours', 0)

wps_config = copy.deepcopy(config)
if spinup_hours > 0:
	wps_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	wps_config['custom']['forecast_hours'] += spinup_hours
run_wps(wps_config)

# Spinup.
if spinup_hours > 0:
	spinup_config = copy.deepcopy(config)
	spinup_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	spinup_config['custom']['forecast_hours'] = spinup_hours
	spinup_config['custom']['end_time'] = spinup_config['custom']['start_time'].add(hours=spinup_hours)
	run_spinup(spinup_config)

# Run reference experiment without DA.
run_osse('ref')

# Run sensitivity experiment with DA.
run_osse('sens')

# Merge two met json files.
try:
	json_file = os.path.basename(glob(f'{args.work_root}/ref/met/d01_*.json')[0])
	d1 = json.load(open(f'{args.work_root}/ref/met/{json_file}', 'r'))
	d2 = json.load(open(f'{args.work_root}/sens/met/{json_file}', 'r'))
	json.dump({'ref': d1,'sens': d2}, open(f'{args.work_root}/result.json', 'w'))
except:
	cli.warning('There is no MET result in JSON format.')
